package geppetto.cat.models;


import geppetto.cat.alignments.Alignment;
import geppetto.cat.common.Common;
import geppetto.cat.constrains.ConstrainedProjectionStats;
import geppetto.cat.corpus.BilingualCorpus;
import geppetto.cat.models.stats.AgreeEStepStats;
import geppetto.cat.models.stats.ProjectionStats;

import java.io.BufferedReader;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.FileReader;
import java.io.IOException;
import java.io.PrintStream;




public class AgreementM1Iterative extends AbstractModel {

	public M1 forward;

	public M1 backward;

	double epsilon;

	int iterations;

	boolean _trained = false;

	public AgreementM1Iterative(BilingualCorpus forwardCorpus,
			BilingualCorpus backwardCorpus, double epsilon, int iterations) {
		this.epsilon = epsilon;
		this.iterations = iterations;
		forward = new M1(forwardCorpus);
		backward = new M1(backwardCorpus);
	}

	AgreementM1Iterative() {

	}

	// ///////////////////////////////////Save Model and Load Model
	// Code//////////////////////
	// ///// Save and load models from file
	public void saveModel(String directory) {
		try {
			PrintStream file = new PrintStream(new FileOutputStream(directory
					+ "epsilon"));
			file.println(iterations);
			file.println(epsilon);
			forward.saveModel(directory + "forward");
			backward.saveModel(directory + "backward");
		} catch (FileNotFoundException e) {
			System.out
					.println("Model2 Agrrement Constrain: Could not save model");
			System.exit(-1);
		}
	}

	public static AgreementM1Iterative loadModel(BilingualCorpus corpus,
			String directory) {
		AgreementM1Iterative model = new AgreementM1Iterative();
		try {
			BufferedReader file = new BufferedReader(new FileReader(directory
					+ "epsilon"));
			model.iterations = Integer.parseInt(file.readLine());
			model.epsilon = Double.parseDouble(file.readLine());
			System.out.println("Reading Forward: " + directory + "forward");
			model.forward = M1.loadModel(corpus, directory + "forward");
			System.out.println("Reading Forward: " + directory + "backward");
			model.backward = M1.loadModel(corpus.reverse(), directory
					+ "backward");
		} catch (FileNotFoundException e) {
			System.out
					.println("Model1 Agreement Constrain: Could not read model file not found "
							+ e.toString());
			System.exit(-1);
		} catch (NumberFormatException e) {
			System.out
					.println("Model1 Agreement Constrain: Could not read model number conversion exception"
							+ e.toString());
			System.exit(-1);

		} catch (IOException e) {
			System.out
					.println("Model1 Agreement Constrain: Could not read model"
							+ e.toString());
			System.exit(-1);
		}
		return model;
	}

	public String getName() {
		return "Model 1 with Agreement Training";
	}

	// Calcultes the projection statistics for just one instance
	public ProjectionStats projectPosteriors(double[][] qForward,
			double[][] qBackward, int sSize, int fSize) {
		ProjectionStats stats = new ProjectionStats(iterations);

		// Save original values
		double[][] posteriorForward = geppetto.cat.common.StaticTools.copyMatrix(
				qForward, sSize + 1, fSize);
		double[][] posteriorBackward = geppetto.cat.common.StaticTools.copyMatrix(
				qBackward, fSize + 1, sSize);

		double phi[][] = new double[sSize][fSize];
		for (int grI = 0; grI < iterations; grI++) {
			double lRate = 10.0 / (10.0 + grI);
			double totalViolation = 0;
			for (int fi = 0; fi < fSize; fi++) {
				for (int si = 0; si < sSize; si++) {
					// Update phi and lambda
					double ePsi = qForward[si][fi] - qBackward[fi][si];
					double phiUpdate = lRate * (-ePsi);
					totalViolation += Math.abs(-ePsi);
					double newlambda = Math.abs(phi[si][fi]) - lRate * epsilon;
					double newPhi = phiUpdate + phi[si][fi];

					// Project
					if (-newPhi <= newlambda && newlambda <= newPhi) {
						newPhi = (newPhi + newlambda) / 2;
						newlambda = newPhi;
					} else if (newPhi <= newlambda && newlambda <= -newPhi) {
						newPhi = (newPhi - newlambda) / 2;
						newlambda = -newPhi;
					} else {
						newPhi = 0;
						newlambda = 0;
					}
					qForward[si][fi] = qForward[si][fi] * Math.exp(newPhi);
					qBackward[fi][si] = qBackward[fi][si] * Math.exp(-newPhi);
					phi[si][fi] = newPhi;
				}
			}

			// Renormalizing
			qForward = Common.normalize(qForward, fSize, sSize + 1);
			qBackward = Common.normalize(qBackward, sSize, fSize + 1);

			// Adding Stats
			stats.addViolation(totalViolation, grI);
			// Calculate thw KL(q|p) which is what we are minimizing
			stats.addForwardKL(geppetto.cat.common.StaticTools.KLDistance(qForward,
					posteriorForward, sSize + 1, fSize), grI);
			stats.addBackwardKL(geppetto.cat.common.StaticTools.KLDistance(qBackward,
					posteriorBackward, fSize + 1, sSize), grI);
		}
		return stats;
	}

	public AgreeEStepStats eStep() {
		double totalLikelihoodF = 0;
		double totalLikelihoodB = 0;
		int numZeroLikelihood = 0;
		forward.clearCounts();
		backward.clearCounts();

		AgreeEStepStats d = new AgreeEStepStats();
		ProjectionStats totalPStats = new ProjectionStats(iterations);
		for (int i = 0; i < forward._corpus.getNumberOfTrainingSentences(); i++) {
			int[] s = forward._corpus.getSourceSentence(i,
					BilingualCorpus.TRAIN_CORPUS);
			final int sSize = s.length;
			int[] f = forward._corpus.getForeignSentence(i,
					BilingualCorpus.TRAIN_CORPUS);
			final int fSize = f.length;
			double[][] forwardProbCache = forward.makeProbCache(f, s);
			double[][] backwardProbCache = backward.makeProbCache(s, f);
			// indexed by[si][fi] sSize corresponds to null word probability
			double[][] posteriorsF = new double[sSize + 1][fSize];
			totalLikelihoodF += forward.calculatePosteriors(s, f, posteriorsF,forwardProbCache);
			// indexed by[fi][si] fSize corresponds to null word probability
			double[][] posteriorsB = new double[fSize + 1][sSize];
			totalLikelihoodB += backward.calculatePosteriors(f, s, posteriorsB, backwardProbCache);
			totalPStats.add(projectPosteriors(posteriorsF, posteriorsB, sSize,
					fSize));
			forward.addCounts(s, f, posteriorsF);
			backward.addCounts(f, s, posteriorsB);
		}

		if (numZeroLikelihood > 0) {
			System.out.println("Numerical underflow happened "
					+ numZeroLikelihood + " times");
		}
		d.logLikelihoodF = totalLikelihoodF;
		d.logLikelihoodB = totalLikelihoodB;
		d.numSents = forward._corpus.getNumberOfTrainingSentences();
		d.pStats = totalPStats;
		return d;
	}

	public void setTrained() {
		_trained = true;
		forward.setTrained();
		backward.setTrained();
	}

	public void mStep() {
		forward.mStep();
		backward.mStep();
	}

	public Alignment posteriorDecodingAlignment(int sentenceNumber,
			byte sentenceSource, float treshhold,boolean projectPosteriors, ConstrainedProjectionStats stats) {
		return forward.posteriorDecodingAlignment(sentenceNumber,
				sentenceSource, treshhold,false,null);
	}

	public Alignment viterbiAlignment(int sentenceNumber, byte sentenceSource, boolean projectPosteriors, ConstrainedProjectionStats stats) {
		if(projectPosteriors){
			System.out.println("Ignoring project posteriors");
		}
		return forward.viterbiAlignment(sentenceNumber, sentenceSource, projectPosteriors,null);
	}

	

	
	public void initializeTrain() {
		forward.initializeTrain();
		backward.initializeTrain();
	}

	public void finalizeTrain() {
		forward.finalizeTrain();
		backward.finalizeTrain();
	}

	// TODO fornow just use the forward model
	public double getNullPhrasePosterior(int phraseNumber, byte phraseSource,
			int[] foreingSentence, int[] sourceSentence, int startForeignIndex,
			int endForeignIndex) {
		return forward.getNullPhrasePosterior(phraseNumber, phraseSource,
				foreingSentence, sourceSentence, startForeignIndex,
				endForeignIndex);
	}

	// TODO fornow just use the forward model
	public double getPhrasePosterior(int phraseNumber, byte phraseSource,
			int[] foreingSentence, int[] sourceSentence, int startSourceIndex,
			int endSourceIndex, int startForeignIndex, int endForeignIndex) {
		return forward.getPhrasePosterior(phraseNumber, phraseSource,
				foreingSentence, sourceSentence, startSourceIndex,
				endSourceIndex, startForeignIndex, endForeignIndex);
	}

	


}
